package net.tp.algo.trie;

import java.util.ArrayList;
import java.util.List;
import java.util.Stack;

/**
 * 
 * Symbol table with key as String. It's implemented as R-way trie.
 * 
 * @author Trung Phan
 *
 */
public class RwayTrie<V> implements Trie<V> {
	
	private final static int R = 256;
	
	private Node root = new Node();
	
	private int size;
	
	@Override
	public int size() {
		return this.size;
	}
	
	@Override
	public boolean empty() {
		return size == 0;
	}
	
	@Override
	public void put(String key, V value) {
		if (value == null || key == null || key.length() == 0) {
			return;
		}

		Node n = root;
		for (int i = 0; i < key.length(); i++) {
			char c = key.charAt(i);
			if (n.next[c] == null) {
				n.next[c] = new Node();
			}
			n = n.next[c];
		}
		
		if (n.value == null) {
			size++;
		}
		n.value = value;
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public V get(String key) {
		Node n = findNode(key);
		return n == null ? null : (V)n.value;
	}
	
	private Node findNode(String key) {
		Node n = root;
		for (int i = 0; i < key.length(); i++) {
			char c = key.charAt(i);
			n = n.next[c];
			if (n == null) {
				return null;
			}
		}
		return n;
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public V delete(String key) {
		// stack of visited nodes to remove later if empty.
		Stack<Node> visitedNodes = new Stack<>();
		
		Node n = root;
		for (int i = 0; i < key.length(); i++) {
			visitedNodes.push(n);
			char c = key.charAt(i);
			n = n.next[c];
			if (n == null) {
				return null;
			}
		}
		
		if (n.value != null) {
			size--;
		}
		
		V result = (V)n.value;
		n.value = null;
		
		// remove empty nodes
		for (int i = key.length() - 1; i >= 0 && !visitedNodes.empty(); i--) {
			n = visitedNodes.pop();
			char c = key.charAt(i);
			if (n.next[c].isEmpty()) {
				n.next[c] = null;
			}
		}
		
		return result;
	}
	
	/**
	 * Doing depth first search
	 * @return keys in dfs visit order
	 */
	@Override
	public Iterable<String> keys() {
		
		final List<String> result = new ArrayList<>();
		
		dfs("", root, new Visitor<V>(){
			public void visit(String key, V value) {
				result.add(key);
			}
		});
		
		return result;
		
	}
	
	@Override
	public Iterable<String> keysWithPrefix(String s) {
		
		final List<String> result = new ArrayList<>();
		
		Node n = findNode(s);
		if (n != null) {
			if (n.value != null) {
				result.add(s);
			}
			
			dfs(s, n, new Visitor<V>() {
				public void visit(String key, V value) {
					result.add(key);
				}
			});
		}
		
		
		return result;
		
	}
	
	/**
	 * @see Trie#longestPrefixOf(String)
	 */
	@Override
	public String longestPrefixOf(String s) {
		
		String result = "";
		
		StringBuilder prefix = new StringBuilder();
		Node n = root;
		for (int i = 0; i < s.length(); i++) {
			char c = s.charAt(i);
			if (n.next[c] == null) {
				break;
			}
			prefix.append(c);
			n = n.next[c];
			if (n.value != null) {
				result = prefix.toString();
			}
		}
		
		return result;
	}
	
	/**
	 * keys that match s (where . is a wildcard)
	 * @param s string
	 * @return keys
	 */
	public Iterable<String> keysThatMatch(String s) {
		
		List<String> result = new ArrayList<>();
		
		if (root == null) {
			return result;
		}
		Stack<StackFrame> stack = new Stack<>();
		stack.push(new StackFrame(root));
		StringBuilder keyBuilder = new StringBuilder();
		
		while (!stack.empty()) {
			StackFrame f = stack.peek();
			Node n = f.node;
			
			char searchC = s.charAt(stack.size() - 1);
			if (searchC == '.') {
				while (++f.index < R && n.next[f.index] == null);
			}
			else {
				f.index = f.index == (char)-1 ? searchC : R;
			}
			
			if (f.index == R) {
				stack.pop();
				if (keyBuilder.length() > 0) {
					keyBuilder.deleteCharAt(keyBuilder.length() - 1);
				}
			}
			else {
				Node next = n.next[f.index];
				keyBuilder.append(f.index);
				if (stack.size() == s.length()) {
					if (next.value != null) {
						result.add(keyBuilder.toString());
					}
					keyBuilder.deleteCharAt(keyBuilder.length() - 1);
				}
				else {
					stack.push(new StackFrame(next));
				}
			}
		}

		
		return result;
		
	}
	
	@SuppressWarnings("unchecked")
	private static <V> void dfs(String keyPrefix, Node root, Visitor<V> visitor) {
		if (root == null) {
			return;
		}
		Stack<StackFrame> stack = new Stack<>();
		stack.push(new StackFrame(root));
		StringBuilder keyBuilder = new StringBuilder(keyPrefix);
		
		while (!stack.empty()) {
			StackFrame f = stack.peek();
			Node n = f.node;
			
			while (++f.index < R && n.next[f.index] == null);
			
			if (f.index == R) {
				stack.pop();
				if (keyBuilder.length() > 0) {
					keyBuilder.deleteCharAt(keyBuilder.length() - 1);
				}
			}
			else {
				Node next = n.next[f.index];
				keyBuilder.append(f.index);
				stack.push(new StackFrame(next));
				if (next.value != null) {
					visitor.visit(keyBuilder.toString(), (V)next.value);
				}
			}
		}
	}
	
	private static interface Visitor<V> {
		public void visit(String key, V value);
	}

	/**
	 * Simulated stack frame to implement depth first search.
	 * @author Trung Phan
	 *
	 */
	private static class StackFrame {
		private final Node node;
		private char index;
		
		public StackFrame(Node node) {
			this.node = node;
			index = (char)-1;
		}
	}
	
	private static class Node {
		
		private Object value;
		
		private Node[] next = new Node[R];
		
		public boolean isEmpty() {
			for (int i = 0; i < R; i++) {
				if (next[i] != null) {
					return false;
				}
			}
			return true;
		}
		
		@Override
		public String toString() {
			return value == null ? "." : value.toString();
		}
	}
	
}
